package com.zhongfu.ai.common.config;

import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;

@Configuration
public class RagEmbeddingConfig {


    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }

    @Bean
    public PgVectorStore pgVectorStore(DynamicRoutingDataSource dynamicRoutingDataSource, OpenAiEmbeddingModel embeddingModel) {
        DataSource dataSource = dynamicRoutingDataSource.getDataSource("pgvector");
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        PgVectorStore.PgVectorStoreBuilder pgVectorStoreBuilder = PgVectorStore.builder(jdbcTemplate, embeddingModel);
        pgVectorStoreBuilder.initializeSchema(true);
        pgVectorStoreBuilder.dimensions(1536);
        return pgVectorStoreBuilder.build();
    }


    /**
     * 如果使用OpenAI模型进行搭建RAG知识库，PgVector需要进行下面的配置
     */

//    @Bean
//    public SimpleVectorStore simpleVectorStore(OpenAiEmbeddingModel embeddingModel){
//        SimpleVectorStore.SimpleVectorStoreBuilder simpleVectorStoreBuilder =  SimpleVectorStore.builder(embeddingModel);
//        return simpleVectorStoreBuilder.build();
//    }
//
//    @Bean
//    public PgVectorStore pgVectorStore(JdbcTemplate jdbcTemplate, OpenAiEmbeddingModel embeddingModel){
//        PgVectorStore.PgVectorStoreBuilder pgVectorStoreBuilder = PgVectorStore.builder(jdbcTemplate, embeddingModel);
//        pgVectorStoreBuilder.initializeSchema(true);
//        pgVectorStoreBuilder.dimensions(768);
//        return pgVectorStoreBuilder.build();
//    }
}
